#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HF 模型/数据集下载工具（自动使用 hf-mirror.com 镜像从huggingface.co下载）
支持自动启用 hf-transfer、智能默认路径、符号链接禁用等。

使用示例：
  python hfdownload.pyw --repo_id speechbrain/spkrec-ecapa-voxceleb
"""

import os
import sys
import argparse
import importlib.util
from pathlib import Path
from huggingface_hub import snapshot_download


# ==============================
# 1. 自动配置环境变量
# ==============================

# 1.1 检查并设置 HF_ENDPOINT 为 hf-mirror.com（注意：去除多余空格）
HF_MIRROR_URL = "https://hf-mirror.com"
current_endpoint = os.environ.get("HF_ENDPOINT", "").strip()

if current_endpoint != HF_MIRROR_URL:
    os.environ["HF_ENDPOINT"] = HF_MIRROR_URL
    print(f"✅ 已设置 HF_ENDPOINT = {HF_MIRROR_URL}")

# 1.2 检查是否安装了 hf-transfer，若安装则启用加速
if importlib.util.find_spec("hf_transfer") is not None:
    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
    print("✅ 已启用 hf-transfer 加速下载")
else:
    print("ℹ️ 未安装 hf-transfer，如需加速可运行: pip install hf-transfer")


# ==============================
# 2. 参数解析
# ==============================
def parse_args():
    parser = argparse.ArgumentParser(
        description="Hugging Face 模型下载工具（自动使用 hf-mirror.com 镜像）",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        "--repo_id",
        type=str,
        required=True,
        help="Hugging Face 仓库 ID，格式：username/repo_name\n"
             "例如：speechbrain/spkrec-ecapa-voxceleb"
    )
    parser.add_argument(
        "--repo_type",
        type=str,
        default="model",
        choices=["model", "dataset", "space"],
        help="仓库类型（默认: model）"
    )
    parser.add_argument(
        "--local_dir",
        type=str,
        default=None,
        help="本地保存目录（默认：./models/<repo_id 转换后路径>）"
    )
    parser.add_argument(
        "--allow_patterns",
        type=str,
        default=None,
        help="仅下载匹配的文件（如 '*.safetensors,tokenizer.json'），逗号分隔"
    )
    parser.add_argument(
        "--ignore_patterns",
        type=str,
        default=None,
        help="跳过匹配的文件（如 '*.bin'），逗号分隔"
    )
    parser.add_argument(
        "--token",
        type=str,
        default=None,
        help="Hugging Face Token（用于私有或 Gated 仓库）"
    )
    parser.add_argument(
        "--revision",
        type=str,
        default=None,
        help="Git 分支或 commit hash（默认：main）"
    )
    return parser.parse_args()


# ==============================
# 3. 主下载逻辑
# ==============================
def main():
    args = parse_args()

    # 3.1 处理 local_dir：若未指定，默认为 ./models/<repo_id 替换 '/' 为系统路径分隔符>
    if args.local_dir is None:
        # 将 repo_id 中的 '/' 转为合法路径（如 Windows 用 \，Linux/macOS 用 /）
        safe_repo_path = args.repo_id.replace("/", os.sep)
        local_dir = Path("models") / safe_repo_path
        local_dir = local_dir.resolve()
    else:
        local_dir = Path(args.local_dir).resolve()

    # 3.2 处理文件过滤模式
    allow_patterns = args.allow_patterns.split(",") if args.allow_patterns else None
    ignore_patterns = args.ignore_patterns.split(",") if args.ignore_patterns else None

    # 3.3 打印信息
    print(f"\n📥 准备下载仓库: {args.repo_id}")
    print(f"📁 保存到: {local_dir}")
    print(f"🌐 使用镜像: {os.environ['HF_ENDPOINT']}")

    # 3.4 执行下载
    try:
        downloaded_path = snapshot_download(
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,  # ✅ 默认不使用符号链接（所见即所得）
            allow_patterns=allow_patterns,
            ignore_patterns=ignore_patterns,
            token=args.token,
            revision=args.revision,
        )
        print(f"\n✅ 下载成功！路径: {downloaded_path}")
    except Exception as e:
        print(f"\n❌ 下载失败: {e}")
        sys.exit(1)


# ==============================
# 4. 入口
# ==============================
if __name__ == "__main__":
    main()